open Base

include Unique_id.Ids

module Dyn = Dyn
module Src = Src
module Behavior = Behavior
module Event = Event
module Tree = Tree
module Comp = Comp

(****************************************)

module State = struct

  type t =
    { src_vals : Src.values
    ; conns : Sink_id.t list Src_id.map
    ; sinks : sink Sink_id.map }
  and state = t

  and 'a typed_sink =
    { tree : Tree.t
    ; render : Src.values -> 'a Comp.t
    ; env : 'a Comp.env
    ; del : to_del list }
  and sink = Sink : 'a typed_sink -> sink

  and to_add =
    [ `Sink of Sink_id.t * sink
    | `Src of Src_id.t * Src.value
    | `Conn of Src_id.t * Sink_id.t ]

  and to_del =
    [ `Sink of Sink_id.t
    | `Src of Src_id.t ]

  let empty =
    { src_vals = Src.empty_values
    ; conns = Src_id.empty_map
    ; sinks = Sink_id.empty_map }

  let of_sources al =
    { empty
      with src_vals = Map.of_alist_exn (module Src_id) al }

  (* manipulating state *)

  let add t : to_add list -> t =
    List.fold ~init:t ~f:(fun t -> function
        | `Src(src, init) ->
           { t with src_vals = Map.add_exn t.src_vals ~key:src ~data:init }
        | `Sink(id, sink) ->
           { t with sinks = Map.add_exn t.sinks ~key:id ~data:sink }
        | `Conn(src, sink) ->
           { t with conns = Map.add_multi t.conns ~key:src ~data:sink })

  let rec delete t : to_del list -> t =
    List.fold ~init:t
      ~f:(fun t -> function
          | `Sink(sink_id) ->
             delete_sink t sink_id
          | `Src(src) ->
             { t with src_vals = Map.remove t.src_vals src })

  and delete_sink t sink_id =
    match Map.find t.sinks sink_id with
    | None -> t
    | Some(Sink sink) ->
       let not_this_sink id = not (Sink_id.equal id sink_id) in
       let sinks = Map.remove t.sinks sink_id in
       let conns = Map.map t.conns ~f:(List.filter ~f:not_this_sink) in
       delete { t with sinks ; conns }
         sink.del

  (* mounting *)

  type ctx =
    { mutable ctx_src_vals : Src.values
    ; mutable to_add : to_add list
    ; mutable to_del : to_del list }

  include Comp.Make_inst(struct
      type t = ctx

      type deleter =
        to_del list

      let src_vals {ctx_src_vals;_} = ctx_src_vals

      let capture_deleter ctx ~f =
        let to_del = ctx.to_del in
        let () = ctx.to_del <- [] in
        let output = f () in
        let to_del' = ctx.to_del in
        let () = ctx.to_del <- to_del in
        (to_del', output)

      let create_source ctx ~init =
        let src = Src_id.create () in
        let () = ctx.ctx_src_vals <- Map.set ctx.ctx_src_vals ~key:src ~data:init in
        let () = ctx.to_add <- `Src(src, init) :: ctx.to_add in
        let () = ctx.to_del <- `Src(src) :: ctx.to_del in
        src

      let create_sink ctx ~tree ~env ~render ~del =
        let sink_id = Sink_id.create () in
        let sink = Sink { tree ; del ; render ; env } in
        let () = ctx.to_add <- `Sink(sink_id, sink) :: ctx.to_add in
        let () = ctx.to_del <- `Sink(sink_id) :: ctx.to_del in
        sink_id

      let connect ctx sink_id src_id =
        ctx.to_add <- `Conn(src_id, sink_id) :: ctx.to_add
    end)

  let inst t env comp =
    let ctx = { ctx_src_vals = t.src_vals ; to_add = [] ; to_del = [] } in
    let tree = inst ctx env comp in
    (ctx.to_add, ctx.to_del, tree)

  let mount t comp : Tree.t * t =
    let (to_add, _to_del, tree) = inst t Comp.empty_env comp in
    let () = Tree.refresh_path [] tree in
    let t = add t to_add in
    (tree, t)

  let rerender t sink : _ typed_sink * t =
    let comp = sink.render t.src_vals in
    let (to_add', to_del', tree') = inst t sink.env comp in
    let sink' = { sink with tree = tree' ; del = to_del' } in
    let t = delete t sink.del in
    let t = add t to_add' in
    (sink', t)

  let update t effect : (Tree.t * Tree.t) list * t =
    let t = { t with src_vals = Src.apply_effect t.src_vals effect } in
    let (upds_rev, t) =
      Sequence.(
        fold ~init:([], t) ~f:(fun (upds_rev, t) sink_id ->
            match Map.find t.sinks sink_id with
            | None -> (upds_rev, t)
            | Some(Sink sink) ->
               let (sink', t) = rerender t sink in
               let t = { t with sinks = Map.set t.sinks ~key:sink_id ~data:(Sink sink') } in
               let upd = (sink.tree, sink'.tree) in
               (upd::upds_rev, t))
        @@
        (* get the sink id's associated with the affected sources *)
        concat_map (of_list (Map.keys effect))
          ~f:(fun src -> of_list (Map.find_multi t.conns src))
      )
    in
    (List.rev upds_rev, t)

  let update1 t src fn =
    update t (Src.single_effect src fn)

end
